import numpy as np
from image_process.ellipse_fit import ellipse_fit_score, fit_ellipse

def try_roi_process(pos_img, neg_img, ellipse, scale=1.25, min_score=20):
    """
    根据上一帧椭圆参数生成ROI掩码, 并在ROI内直接拟合椭圆
    :param pos_img: 正极性事件计数图 (H, W)
    :param neg_img: 负极性事件计数图 (H, W)
    :param ellipse: (center(x, y), (width, height), angle)
    :param scale: ROI放大倍数
    :param min_points: 拟合椭圆的最小点数
    :return: ellipse(None表示拟合失败)
    """
    shape = pos_img.shape
    (x_center, y_center), (width, height), angle = ellipse
    width *= scale
    height *= scale
    Y, X = np.ogrid[:shape[0], :shape[1]]
    angle_rad = np.deg2rad(-angle)
    x_shift = X - x_center
    y_shift = Y - y_center
    x_rot = x_shift * np.cos(angle_rad) - y_shift * np.sin(angle_rad)
    y_rot = x_shift * np.sin(angle_rad) + y_shift * np.cos(angle_rad)
    mask = ((x_rot / (width/2))**2 + (y_rot / (height/2))**2) <= 1
    roi_mask = mask.astype(np.uint8)

    # 只保留ROI内的事件
    event_img = (pos_img + neg_img) * roi_mask
    points_y, points_x = np.where(event_img > 0)
    points = np.column_stack((points_x, points_y))
    
    # 拟合椭圆
    new_ellipse = fit_ellipse(points)
    
    # 检查拟合分数
    fit_score = 0

    if new_ellipse is not None:
        fit_score = ellipse_fit_score(points, new_ellipse)

    if fit_score > min_score:
        return new_ellipse
    else:
        return None